package com.essence.common.task;

import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.support.CronTrigger;

import java.util.Date;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

public class TaskLauncher {
    private final ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
    private Map<String, TaskExecutor> taskTable = new ConcurrentHashMap<>();

    public TaskLauncher(int poolSize) {
        scheduler.setThreadNamePrefix("task-");
        scheduler.setPoolSize(poolSize);
        scheduler.setRemoveOnCancelPolicy(true);
        scheduler.setWaitForTasksToCompleteOnShutdown(true);
        scheduler.setAwaitTerminationSeconds(60);
        scheduler.initialize();
    }

    public void addTask(TaskExecutor task) {
        task.setTaskLauncher(this);
        if (task.taskType() == TaskExecutor.TaskType.INTERVAL) {
            Date startTime = new Date(System.currentTimeMillis() + ((IntervalTask) task).getDelaySecond());
            ScheduledFuture<?> scheduledFuture = scheduler.scheduleAtFixedRate(task, startTime, task.getSecond());
            task.setFuture(scheduledFuture);
        }
        if (task.taskType() == TaskExecutor.TaskType.DELAYED) {
            Date startTime = new Date(System.currentTimeMillis() + ((DelayedTask) task).getDelaySecond());
            ScheduledFuture<?> scheduledFuture = scheduler.scheduleWithFixedDelay(task, startTime, task.getSecond());
            task.setFuture(scheduledFuture);
        }
        if (task.taskType() == TaskExecutor.TaskType.ONCE) {
            Date startTime = new Date(System.currentTimeMillis() + ((OnceTask) task).getDelaySecond());
            ScheduledFuture<?> scheduledFuture = scheduler.scheduleAtFixedRate(task, startTime, task.getSecond());
            task.setFuture(scheduledFuture);
        }
        if (task.taskType() == TaskExecutor.TaskType.CRON) {
            ScheduledFuture<?> scheduledFuture = scheduler.schedule(task, new CronTrigger(task.getCronExpression()));
            task.setFuture(scheduledFuture);
        }
        taskTable.put(task.getTaskName(), task);
    }

    public boolean removeTask(String taskName) {
        TaskExecutor taskExecutor = taskTable.remove(taskName);
        if (taskExecutor == null) {
            return false;
        }
        if (taskExecutor.cancel(false)) {
            return scheduler.getScheduledThreadPoolExecutor().remove(taskExecutor);
        }
        return false;
    }

    public Set<String> taskNames() {
        return taskTable.keySet();
    }

    public TaskExecutor getTaskExecutor(String taskName) {
        return taskTable.get(taskName);
    }
}
